#!/usr/bin/env bash

script_version="0.1.0"
echo Script Version: $script_version
line_thing="------------------------------------"
echo "$line_thing"
echo "Welcome to the Vast.ai Ai horde setup script! Please follow the on-screen instructions to continue."
echo "$line_thing"
amount_of_gpus=$(nvidia-smi --query-gpu=count --format=csv,noheader | head -n 1)
amount_of_ram=$(free -h | awk '/^Mem:/ {print $2}' | sed 's/[^0-9]*//g')
cpu_cores_fixed=$(nproc)

#default variables
num_gpus=""
apikey=""
queue_size="1"  
max_batch="1" 
safety_on_gpu="false"  
max_power="8"  
high_memory_mode="false"  
high_performance_mode="false"  
moderate_performance_mode="false"  
unload_models_from_vram_often="true"  
allow_lora="false"  
max_lora_cache_size="10"  
models_to_load=('top 2')
load_large_models="false"  
clientname="Vast_instance"
bypass_confirm="n"
autosettings=""
fastdisk='false'
while [[ $# -gt 0 ]]; do
  key="$1"
  case $key in
    -apikey)
      apikey="$2"
      shift 2;;
    -queue_size)
      queue_size="$2"
      shift 2;;
    -num_gpus)
      num_gpus="$2"
      shift 2;;
    -fastdisk)
      fastdisk="$2"
      shift 2;;
    -max_batch)
      max_batch="$2"
      shift 2;;
    -safety_on_gpu)
        if [ "$2" == "y" ]; then
            safety_on_gpu="true"
        else
            safety_on_gpu="false"
        fi
        shift 2;;
    -max_power)
      max_power="$2"
      shift 2;;
    -high_memory_mode)
        if [ "$2" == "y" ]; then
            high_memory_mode="true"
        else
            high_memory_mode="false"
        fi
        shift 2;;
    -high_performance_mode)
        if [ "$2" == "y" ]; then
            high_performance_mode="true"
        else
            high_performance_mode="false"
        fi
        shift 2;;
    -moderate_performance_mode)
        if [ "$2" == "y" ]; then
            moderate_performance_mode="true"
        else
            moderate_performance_mode="false"
        fi
        shift 2;;
    -unload_models_from_vram_often)
        if [ "$2" == "y" ]; then
            unload_models_from_vram_often="true"
        else
            unload_models_from_vram_often="false"
        fi
        shift 2;;
    -allow_lora)
        if [ "$2" == "y" ]; then
            allow_lora="true"
        else
            allow_lora="false"
        fi
        shift 2;;
    -max_lora_cache_size)
      max_lora_cache_size="$2"
      shift 2;;
    -clientname)
      clientname="$2"
      shift 2;;
    -y)
      bypass_confirm="y"
      shift 1;;
    -models)
      IFS=',' read -r -a models_to_load <<< "$2"
      shift 2;;
    -load_large_models)
        if [ "$2" == "y" ]; then
            load_large_models="true"
        else
            load_large_models="false"
        fi
        shift 2;;
    -autosettings)
      autosettings="$2"
      shift 2;;
    *)
    
      echo "Unknown option: $1"
      exit 1
      ;;
  esac
done


if [ -z "$apikey" ]; then 
  read -r -p "Enter your API key: (leave it blank for anonymous) " apikey
  if [ -z "$apikey" ]; then
    apikey="0000000000"
  fi
fi
if [ -z "$autosettings" ]; then
  read -r -p "Do you want to use auto settings? (y/n): " autosettings
fi

if [[ "$autosettings" =~ ^[Yy]$ ]]; then
    echo "$line_thing"
    echo "Auto settings enabled."
    num_gpus=$amount_of_gpus
else
    echo "$line_thing"
    echo "There are $amount_of_gpus NVIDIA GPUs in this machine."
    while true; do
      read -r -p "Enter the number of GPUs you would like to use for the AI Horde (1-$amount_of_gpus): " num_gpus < /dev/tty
      if [[ "$num_gpus" =~ ^[1-9][0-9]*$ && "$num_gpus" -le "$amount_of_gpus" ]]; then
        break
      else
        echo "Please enter a valid positive integer for the number of GPUs (1-$amount_of_gpus)."
      fi
    done
fi

echo "$line_thing"
echo "These are the settings selected, is this ok? (Y/n)"
echo "Number of GPUs: $num_gpus"
echo "Client name: $clientname"
echo "API Key: $apikey"
echo "Auto settings: $autosettings"
echo "Safety on GPU: $safety_on_gpu"
echo "Queue Size: $queue_size"
echo "Max Batch: $max_batch"
echo "Max Power: $max_power"
echo "High Memory Mode: $high_memory_mode"
echo "High Performance Mode: $high_performance_mode"
echo "Moderate Performance Mode: $moderate_performance_mode"
echo "Unload models from VRAM often: $unload_models_from_vram_often"
echo "Allow Lora: $allow_lora"
echo "Max Lora Cache Size: $max_lora_cache_size"
echo "Models to load: ${models_to_load[@]}"
echo "Load large models: $load_large_models"
echo "Fast Disk Mode: $fastdisk"


if [[ "$bypass_confirm" != "y" ]]; then
  read -r -p "(Y/n): " confirm < /dev/tty
  if [[ "$confirm" =~ ^[Nn]$ ]]; then
      echo "Aborted"
      exit 0
  fi
else
  echo Bypass on, continuing. 
fi

sudo apt update && sudo apt upgrade -y
sudo apt install bzip2 wget sudo libjemalloc2 git libgl1 tmux -y
sudo mkdir /horde && cd /horde 
sudo chown -R "$(whoami):$(id -gn)" /horde
if [ ! -d "horde-worker-reGen" ]; then
    git clone https://github.com/Haidra-Org/horde-worker-reGen.git
else
    echo "Directory already exists, skipping git..."
fi

cd horde-worker-reGen
cp bridgeData_template.yaml bridgeData.yaml
tmux new-session -d -s horde
models_to_load_str=$(IFS=,; echo "${models_to_load[*]}")
awk -v new_api_key="$apikey" \
    -v new_queue_size="$queue_size" \
    -v new_max_batch="$max_batch" \
    -v new_safety_on_gpu="$safety_on_gpu" \
    -v new_max_power="$max_power" \
    -v new_high_memory_mode="$high_memory_mode" \
    -v new_high_performance_mode="$high_performance_mode" \
    -v new_moderate_performance_mode="$moderate_performance_mode" \
    -v new_unload_models_from_vram_often="$unload_models_from_vram_often" \
    -v new_allow_lora="$allow_lora" \
    -v new_max_lora_cache_size="$max_lora_cache_size" \
    -v new_models_to_load="$models_to_load_str" \
    -v new_load_large_models="$load_large_models" \
    -v new_fastdisk="$fastdisk" '
    BEGIN { skip_models=0 }
    /^[[:space:]]*#/ { print; next }

    /api_key: "0000000000"/ { print "api_key: \"" new_api_key "\""; next }
    /queue_size: [0-9]+/ { print "queue_size: " new_queue_size; next }
    /max_batch: [0-9]+/ { print "max_batch: " new_max_batch; next }
    /safety_on_gpu: [a-zA-Z]+/ { print "safety_on_gpu: " new_safety_on_gpu; next }
    /max_power: [0-9]+/ { print "max_power: " new_max_power; next }
    /high_memory_mode: [a-zA-Z]+/ { print "high_memory_mode: " new_high_memory_mode; next }
    /high_performance_mode: [a-zA-Z]+/ { print "high_performance_mode: " new_high_performance_mode; next }
    /moderate_performance_mode: [a-zA-Z]+/ { print "moderate_performance_mode: " new_moderate_performance_mode; next }
    /unload_models_from_vram_often: [a-zA-Z]+/ { print "unload_models_from_vram_often: " new_unload_models_from_vram_often; next }
    /allow_lora: [a-zA-Z]+/ { print "allow_lora: " new_allow_lora; next }
    /max_lora_cache_size: [0-9]+/ { print "max_lora_cache_size: " new_max_lora_cache_size; next }
    /models_to_load:/ {
        print "models_to_load:";
        n = split(new_models_to_load, models, ",");
        for (i = 1; i <= n; i++) {
            # Trim any extra whitespace from each model
            gsub(/^[ \t]+|[ \t]+$/, "", models[i]);
            if (length(models[i]) > 0)
                print "    - \"" models[i] "\"";
        }
        skip_models=1;
        next;
    }
    skip_models && $0 ~ /^[[:space:]]+/ { next }
    skip_models && $0 !~ /^[[:space:]]+/ { skip_models=0 }
    /load_large_models: [a-zA-Z]+/ { print "load_large_models: " new_load_large_models; next }
    /fastdisk: [a-zA-Z]+/ { print "fastdisk: " new_fastdisk; next }
    { print }
' bridgeData.yaml > tmp.yaml && mv tmp.yaml bridgeData.yaml


cd /horde/horde-worker-reGen

sh preload-models.sh
echo Models loaded, starting clients.

for (( i = 0; i < num_gpus; i++ )); do
    HOSTNAME="${clientname}-$(printf "%03d" $(( i + 1 )))"
    
    CMD="cd /horde/horde-worker-reGen && CUDA_VISIBLE_DEVICES=$i ./horde-bridge.sh -n $HOSTNAME-$i"
    
    if [ "$i" -eq 0 ]; then
        tmux rename-window -t horde "gpu-$i"
        tmux send-keys -t horde "$CMD" C-m
    else
        tmux new-window -t horde -n "gpu-$i"
        tmux send-keys -t horde:"gpu-$i" "$CMD" C-m
    fi
done

tmux switch-client -t horde
tmux attach-session -t horde
echo "If it doesn't switch automatically, you can switch with 'tmux attach-session -t horde' or 'tmux switch-client -t horde'"